#encoding=utf-8
import rospy
import cv2
from sensor_msgs.msg import Image
from rosgraph_msgs.msg import Log
from cv_bridge import CvBridge, CvBridgeError
import numpy as np 
import message_filters
from copy import deepcopy

import sys, select, termios, tty
from commander1 import Commander
import time
import threading

def getKey():
	tty.setraw(sys.stdin.fileno())
	select.select([sys.stdin], [], [], 0)
	key = sys.stdin.read(1)
	termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
	return key

class ObstacleAvoidance:
	def __init__(self, depth_topic = '/camera/depth/image_raw', rosout_topic = '/rosout', debug = False):
		self.con = Commander()
		self.bridge = CvBridge()
		
		self.pred_time = .3 # 向前预测的时间
		self.velocity_threshold = 3 # 速度阈值，单位 m/s，当大于这个阈值时认为是在运动的点
		self.move_threshold = 0 # 决策阈值，当预测的深度图像有小于这个值的时候，开始躲避
		self.direction_ratio = 20 # 当比例超过这个值时左右，否则上下
		
		self.depth_last_frame = None
		self.last_frame_time = None
		self.main_left_right_move_velocity = 1.5
		self.main_up_move_velocity = 2.
		self.main_moving = False # 是否正在做躲避动作
		self.key_horizontal_move_velocity = 2. # 键盘控制前后左右飞行速度
		self.key_vertical_move_velocity = .5 # 键盘控制升降速度
		self.key_turn_velocity = 20. # 键盘控制自转角速度
		self.turn_angle = 0. # 飞机自转角度
		
		self.debug = debug
		self.debug_file = None
		if debug:
			self.debug_file = open("data", "w")
		
		rospy.Subscriber(depth_topic, Image, self.main)
		rospy_thread = threading.Thread(name = 'rospy_thread', target = rospy.spin)
		rospy_thread.start()
		self.keyControl()
	
	# '''
	def keyControl(self):
		while True:
			if self.keyOnce() is not None: break
	def keyOnce(self):
		if self.main_moving: return None
		
		cur_key = getKey()
		if cur_key == 'x': return True
		keys = {'w' : (1, 0, 0), 'a' : (0, 1, 0), 's' : (-1, 0, 0), 'd' : (0, -1, 0), 'i' : (0, 0, 1), 'k' : (0, 0, -1), 'q' : None, 'e' : None}
		state = dict((x, x == cur_key) for x in "wasdikqe")
		
		move_vec = np.array([0., 0., 0.], dtype = float)
		for key in "wasdik":
			if state[key]: move_vec += np.array(keys[key])
		if move_vec.any():
			self.con.move(move_vec[0] * self.key_horizontal_move_velocity,\
				move_vec[1] * self.key_horizontal_move_velocity,\
				move_vec[2] * self.key_vertical_move_velocity)
			return None
		
		if state['q']:
			self.turn_angle += self.key_turn_velocity
			self.con.turn(self.turn_angle)
		elif state['e']:
			self.turn_angle -= self.key_turn_velocity
			self.con.turn(self.turn_angle)
		
		return None
	# '''
	
	def main(self, data):
		img = self.bridge.imgmsg_to_cv2(data)
		cur_time = time.time()
		depth_img = deepcopy(img)
		depth_img[np.isnan(depth_img)] = 3.
		depth_img[depth_img >= 2.3] = 3.
		
		if self.last_frame_time is not None:
			assert self.depth_last_frame is not None
			delta, delta_time = depth_img - self.depth_last_frame, cur_time - self.last_frame_time
			predict = depth_img + delta * self.pred_time / delta_time
			# self.pred_time 秒后的深度图像预测（直接线性预测）
			h, w = predict.shape
			l_predict, r_predict = predict[:, :w/2], predict[:, w/2+1:]
			l_delta, r_delta = delta[:, :w/2], delta[:, w/2+1:]
			l_fit = ((l_predict <= self.move_threshold) * (l_delta <= -self.velocity_threshold * delta_time)).sum()
			r_fit = ((r_predict <= self.move_threshold) * (r_delta <= -self.velocity_threshold * delta_time)).sum()
			# print("{}, {}".format(l_fit, r_fit))
			if ((predict <= self.move_threshold) * (delta <= -self.velocity_threshold * delta_time)).any():
				self.main_moving = True
				if l_fit > r_fit * self.direction_ratio: # fly to right
					self.con.move(0, -self.main_left_right_move_velocity, 0)
					print('move right!')
				elif r_fit > l_fit * self.direction_ratio: # fly to left
					self.con.move(0, self.main_left_right_move_velocity, 0)
					print('move left!')
				else: # fly up
					self.con.move(0, 0, self.main_up_move_velocity)
					print('move up!')
				time.sleep(.1)
				self.main_moving = False
			
			if self.debug:
				print('debug!')
				print('{}\n{}'.format(delta_time, delta))
				print('printed!')
				
				while True: pass
		
		self.depth_last_frame = depth_img
		self.last_frame_time = cur_time
		time.sleep(.01)

if __name__ == '__main__':
	settings = termios.tcgetattr(sys.stdin)
	ObstacleAvoidance(debug = False)

